{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8fcff912",
   "metadata": {},
   "source": [
    "# Advanced Builder settings\n",
    "\n",
    "<img align=\"left\" src=\"../end2end_example/cybersecurity/finn-example.png\" alt=\"drawing\" style=\"margin-right: 20px\" width=\"250\"/>\n",
    "\n",
    "In this notebook, we'll use the FINN compiler to generate an FPGA accelerator with a streaming dataflow architecture from a small convolutional network trained on CIFAR-10. The key idea in streaming dataflow architectures is to parallelize across layers as well as within layers by dedicating a proportionate amount of compute resources to each layer, illustrated on the figure to the left. You can read more about the general concept in the [FINN](https://arxiv.org/pdf/1612.07119) and [FINN-R](https://dl.acm.org/doi/pdf/10.1145/3242897) papers. This is done by mapping each layer to a Vitis HLS or RTL description, parallelizing each layer's implementation to the appropriate degree and using on-chip FIFOs to link up the layers to create the full accelerator.\n",
    "These implementations offer a good balance of performance and flexibility, but building them by hand is difficult and time-consuming. This is where the FINN compiler comes in: it can build streaming dataflow accelerators from an ONNX description to match the desired throughput."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a830e730",
   "metadata": {},
   "source": [
    "In this tutorial, we will have a more detailed look into the FINN builder tool and explore different options to customize your FINN design. We assume that you have already completed the [Cybersecurity notebooks](../end2end_example/cybersecurity) and that you have a basic understanding of how the FINN compiler works and how to use the FINN builder tool."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ec9a0db",
   "metadata": {},
   "source": [
    "## Outline\n",
    "---------------\n",
    "\n",
    "1. [Introduction to the CNV-w2a2 network](#intro_cnv)\n",
    "2. [Recap default builder flow](#recap_builder)\n",
    "3. [Build steps](#build_step)\n",
    "    1. [How to create a custom build step](#custom_step)\n",
    "4. [Specialize layers configuration json](#specialize_layers)\n",
    "5. [Folding configuration json](#folding_config)\n",
    "6. [Additional builder arguments](#builder_arg)\n",
    "    1. [Verification steps](#verify)\n",
    "    2. [Other builder arguments](#other_args)\n",
    "    3. [Examples for additional builder arguments & bitfile generation](#example_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dbed63f",
   "metadata": {},
   "source": [
    "## Introduction to the CNV-w2a2 network <a id='intro_cnv'></a>\n",
    "\n",
    "The particular quantized neural network (QNN) we will be targeting in this notebook is referred to as CNV-w2a2 and it classifies 32x32 RGB images into one of ten CIFAR-10 classes. All weights and activations in this network are quantized to two bit, with the exception of the input (which is RGB with 8 bits per channel) and the final output (which is 32-bit numbers). It is similar to the convolutional neural network used in the [cnv_end2end_example](../end2end_example/bnn-pynq/cnv_end2end_example.ipynb) Jupyter notebook.\n",
    "\n",
    "\n",
    "You'll have a chance to interactively examine the layers that make up the network in Netron. We start by setting the build directory to the directory this notebook is in and importing helper functions to use in the notebook to examine ONNX graphs and source code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce459f3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from finn.util.visualization import showInNetron, showSrc\n",
    "import os\n",
    "    \n",
    "build_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fc6444c",
   "metadata": {},
   "source": [
    "In the next step, we will export the trained network directly from Brevitas to the QONNX format. QONNX is the intermediate representation (IR) that is used as the frontend to the FINN compiler. Please note that the internal representation of the network is still the FINN-ONNX format. [QONNX and FINN-ONNX](https://finn.readthedocs.io/en/latest/internals.html#intermediate-representation-qonnx-and-finn-onnx) are extensions to the ONNX format to represent quantization, especially below 8 bit, in ONNX graphs. The main difference is that quantization in QONNX graphs is represented using dedicated quantization nodes ([more about QONNX](https://github.com/fastmachinelearning/qonnx)) while the quantization in FINN-ONNX is an annotation attached to the tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe262964",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from finn.util.test import get_test_model_trained\n",
    "from brevitas.export import export_qonnx\n",
    "from qonnx.util.cleanup import cleanup as qonnx_cleanup\n",
    "\n",
    "cnv = get_test_model_trained(\"CNV\", 2, 2)\n",
    "export_onnx_path = build_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "export_qonnx(cnv, torch.randn(1, 3, 32, 32), export_onnx_path)\n",
    "qonnx_cleanup(export_onnx_path, out_file=export_onnx_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d24b632f",
   "metadata": {},
   "source": [
    "After the export, we call a clean up function on the model. This makes sure, that for example all shapes in the network are inferred, constant folding was applied and all tensors and nodes have unique names. In the next step, we can visualize the graph using Netron. When scrolling through the graph, you can see the Quant nodes that indicate the quantization in the network. In the [first step](https://github.com/Xilinx/finn/blob/main/src/finn/builder/build_dataflow_steps.py#L260) of the FINN builder flow, the network gets converted from the QONNX format to the FINN-ONNX format. That means these Quant nodes will not be present in the graph anymore and instead the quantization will be attached as an annotation to the tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87f59da6",
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron(build_dir+\"/end2end_cnv_w2a2_export.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c764ed76",
   "metadata": {},
   "source": [
    "## Quick recap, how to setup up default builder flow for resource estimations <a id='recap_builder'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a26e5418",
   "metadata": {},
   "source": [
    "As a quick recap, let's set up the builder like we have done in the cybersecurity example to get the resource estimates for our example network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9007705a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Quick recap on how to setup the default builder flow for resource estimations\n",
    "\n",
    "import finn.builder.build_dataflow as build\n",
    "import finn.builder.build_dataflow_config as build_cfg\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "model_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\"\n",
    "model_file = model_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "\n",
    "estimates_output_dir = build_dir + \"/output_estimates_only\"\n",
    "\n",
    "#Delete previous run results if exist\n",
    "if os.path.exists(estimates_output_dir):\n",
    "    shutil.rmtree(estimates_output_dir)\n",
    "    print(\"Previous run results deleted!\")\n",
    "\n",
    "\n",
    "cfg_estimates = build.DataflowBuildConfig(\n",
    "    output_dir          = estimates_output_dir,\n",
    "    mvau_wwidth_max     = 80,\n",
    "    target_fps          = 10000,\n",
    "    synth_clk_period_ns = 10.0,\n",
    "    fpga_part           = \"xc7z020clg400-1\",\n",
    "    steps               = build_cfg.estimate_only_dataflow_steps,\n",
    "    generate_outputs=[\n",
    "        build_cfg.DataflowOutputType.ESTIMATE_REPORTS,\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02e4c0f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "build.build_dataflow_cfg(model_file, cfg_estimates);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fa0b9f5",
   "metadata": {},
   "source": [
    "The output directory was created and we can extract information about our model and also how it was processed in the FINN compiler from the generated files. Let's focus on the intermediate models for now. You can find them in the output directory in the folder \"intermediate_models\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05a941ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -t -r {build_dir}/output_estimates_only/intermediate_models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d746eff3",
   "metadata": {},
   "source": [
    "After each FINN builder step, the graph is saved as .onnx file. In the cell above we sort the intermediate models by time in descending order (`ls -t -r`) to visualize the builder flow. As you can see after the conversion to the FINN-ONNX format (`step_qonnx_to_finn`), the graph is prepared by tidy up and streamlining (`step_tidy_up` and `step_streamline`) and then the high level nodes are converted to HW abstraction layers (`step_convert_to_hw`). Then there is a partition created from all layers that were converted to HW layers (`step_create_dataflow_partition`), then we convert each of the HW abstraction layers into an HLS or RTL variant (`step_specialize_layers`). Afterwards optimizations are applied (`step_target_fps_parallelization`, `step_apply_folding_config` and `step_minimize_bit_width`). In the final step of this example we generate resource and performance reports for the network (`step_generate_estimate_reports`). Use the code below to investigate the network after each step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72de8d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_to_investigate = \"step_qonnx_to_finn.onnx\"\n",
    "showInNetron(build_dir+\"/output_estimates_only/intermediate_models/\"+model_to_investigate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bccebd0d",
   "metadata": {},
   "source": [
    "The analysis of these .onnx files can help us identifying points in the flow in which we might need to intervene and provide the compiler with additional information. When investigating the network after the conversion to HW layers, we can see that there are layers that were not converted. We can see this by clicking on the different nodes. HW layers have the module `finn.custom_op.fpgadataflow`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d86463a",
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron(build_dir+\"/output_estimates_only/intermediate_models/step_convert_to_hw.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2719cc09",
   "metadata": {},
   "source": [
    "As you can see in the graph, the first two nodes (a MultiThreshold and Transpose node) and the last two nodes (a Mul and Add node) are not converted into HW layers. FINN currently only converts integer only operations into HW layers, this means only when the input, output & weights are quantized to integer the node will be converted."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff7fa549",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Important notice:</b> We are working on supporting additional data types and this limitation might disappear in the near future.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e6d942e",
   "metadata": {},
   "source": [
    "When we click on the `global_in` in the graph, we can see that the quantization annotation does not contain a data type. If no data type is set and it can not be derived from the preceeding node, the FINN compiler automatically assumes that the data type is floating point. This is why the first node does not get converted into an HW layer, the input is assumed to be floating point."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b8994e6",
   "metadata": {},
   "source": [
    "The solution to the problem depends on the actual data input.\n",
    "1. The data set is quantized and `global_in` is an integer: We set the data type of the tensor `global_in` before passing the model to the FINN compiler using [helper functions of ModelWrapper](https://finn.readthedocs.io/en/latest/internals.html#helper-functions-for-tensors).\n",
    "2. The data set is not quantized: we can either execute the first layer in software (e.g. as part of the Python driver) or we can add a preprocessing step into the graph."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7504dce7",
   "metadata": {},
   "source": [
    "Even though in the example of the CNVw2a2, the inputs are 32x32 RGB images, so the input values are 8 bit (UINT8) \"quantized\", the input to the exported model is floating point. For training in Brevitas, these values were normalized between 0 and 1.0 and so the exported model expects floating point values as input. \n",
    "This means we are in scenario 2. In the next section we will develop a custom step for the FINN builder flow to add preprocessing to our network.\n",
    "\n",
    "But before we move to the next section, let's take a look at the last two nodes in the graph that were not converted to HW layers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9c2696b",
   "metadata": {},
   "source": [
    "We have two nodes at the end of the graph that we were not able to convert: a floating poing scalar multiplication and addition. These operations are \"left-over\" from streamlining and cannot be merged into a succeeding thresholding operation. \n",
    "\n",
    "Our example is a network for image classification, so the output is a vector of 10 values that give a predicition score for each of the classes in the CIFAR-10 data set. If we are only interested in the Top-1 result of the classification, we can add a post-processing step which inserts a TopK node in the graph. \n",
    "\n",
    "Since the last two layers are scalar operations, they have the same influence on all predicition scores in the output vector and we can safely merge them into the TopK node. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fc8fbf5",
   "metadata": {},
   "source": [
    "These pre-processing and post-processing steps are network dependent and we will need to write **custom steps** that can then be executed using the FINN builder tool.\n",
    "\n",
    "In the next section we will first look into how a standard build step inside FINN looks like and then we will write our own custom steps for pre- and post-processing and add them to the builder configuration."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e561a91",
   "metadata": {},
   "source": [
    "## Build steps <a id='build_step'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb18b21d",
   "metadata": {},
   "source": [
    "The following steps are executed when using the `estimates_only`-flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3fe1186",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\n\".join(build_cfg.estimate_only_dataflow_steps))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd3ef987",
   "metadata": {},
   "source": [
    "You can have a closer look at each step by either using the `showSrc()` function or by accessing the doc string."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "313fac18",
   "metadata": {},
   "outputs": [],
   "source": [
    "import finn.builder.build_dataflow_steps as build_dataflow_steps\n",
    "print(build_dataflow_steps.step_tidy_up.__doc__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "029da0da",
   "metadata": {},
   "outputs": [],
   "source": [
    "import finn.builder.build_dataflow_steps as build_dataflow_steps\n",
    "showSrc(build_dataflow_steps.step_tidy_up)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2809f6a7",
   "metadata": {},
   "source": [
    "Each steps gets the model (`model: ModelWrapper`) and the build configuration (`cfg: DataflowBuildConfig`) as input arguments. Then a certain sequence of transformations is applied to the model. In some of the steps, verification can be run to ensure that the applied transformations have not changed the behaviour of the network. In the end the modified model is returned."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9c2c97f",
   "metadata": {},
   "source": [
    "### How to create a custom build step <a id='custom_step'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "537a44e7",
   "metadata": {},
   "source": [
    "When writing our own custom steps, we use the same pattern. See below the code for the pre-processing for the example network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9d43cc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from finn.util.pytorch import ToTensor\n",
    "from qonnx.transformation.merge_onnx_models import MergeONNXModels\n",
    "from qonnx.core.modelwrapper import ModelWrapper\n",
    "from qonnx.core.datatype import DataType\n",
    "\n",
    "def custom_step_add_pre_proc(model: ModelWrapper, cfg: build.DataflowBuildConfig):\n",
    "    ishape = model.get_tensor_shape(model.graph.input[0].name)\n",
    "    # preprocessing: torchvision's ToTensor divides uint8 inputs by 255\n",
    "    preproc = ToTensor()\n",
    "    export_qonnx(preproc, torch.randn(ishape), \"preproc.onnx\", opset_version=11)\n",
    "    preproc_model = ModelWrapper(\"preproc.onnx\")\n",
    "    # set input finn datatype to UINT8\n",
    "    preproc_model.set_tensor_datatype(preproc_model.graph.input[0].name, DataType[\"UINT8\"])\n",
    "    # merge pre-processing onnx model with cnv model (passed as input argument)\n",
    "    model = model.transform(MergeONNXModels(preproc_model))\n",
    "    return model\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a6798aa",
   "metadata": {},
   "source": [
    "In the next step we can modify the builder configuration to execute a custom sequence of builder steps, including the newly implemented pre-processing custom step.\n",
    "\n",
    "For that we create a list `build_steps` which contains next to the standard steps from the `estimate_only` flow, also the new custom step to add the pre-processing. This list then gets passed in the build configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f00b465",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Builder flow with custom step for pre-processing\n",
    "\n",
    "model_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\"\n",
    "model_file = model_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "\n",
    "output_dir = build_dir + \"/output_pre_proc\"\n",
    "\n",
    "#Delete previous run results if exist\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "    print(\"Previous run results deleted!\")\n",
    "\n",
    "build_steps = [\n",
    "    custom_step_add_pre_proc,\n",
    "    \"step_qonnx_to_finn\",\n",
    "    \"step_tidy_up\",\n",
    "    \"step_streamline\",\n",
    "    \"step_convert_to_hw\",\n",
    "    \"step_create_dataflow_partition\",\n",
    "    \"step_specialize_layers\",\n",
    "    \"step_target_fps_parallelization\",\n",
    "    \"step_apply_folding_config\",\n",
    "    \"step_minimize_bit_width\",\n",
    "    \"step_generate_estimate_reports\",\n",
    "]\n",
    "\n",
    "cfg_estimates = build.DataflowBuildConfig(\n",
    "    output_dir          = output_dir,\n",
    "    mvau_wwidth_max     = 80,\n",
    "    target_fps          = 10000,\n",
    "    synth_clk_period_ns = 10.0,\n",
    "    fpga_part           = \"xc7z020clg400-1\",\n",
    "    steps               = build_steps,\n",
    "    generate_outputs=[\n",
    "        build_cfg.DataflowOutputType.ESTIMATE_REPORTS,\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3a2bcea",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "build.build_dataflow_cfg(model_file, cfg_estimates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51b7dbd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -t -r {build_dir}/output_pre_proc/intermediate_models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4690049f",
   "metadata": {},
   "source": [
    "An intermediate .onnx file after the execution of the custom step was automatically created, let's have a look at the graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87e5651e",
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron(build_dir+\"/output_pre_proc/intermediate_models/custom_step_add_pre_proc.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90c6bef9",
   "metadata": {},
   "source": [
    "The graph is in QONNX format and a division by 255 is inserted in the beginning. We can now use the CIFAR-10 images directly as input to the graph and the new `global_in` tensor is UINT8.\n",
    "\n",
    "You can already have a look on how the intermediate models have changed by modifying the code in the cell above. Before we go into more detail, we will add another custom step to insert the post-processing. In this case this means the insertion of a TopK node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c6f1bd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qonnx.transformation.insert_topk import InsertTopK\n",
    "\n",
    "def custom_step_add_post_proc(model: ModelWrapper, cfg: build.DataflowBuildConfig):\n",
    "    model = model.transform(InsertTopK(k=1))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57adbb44",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Builder flow with custom step for pre-processing and post-processing\n",
    "\n",
    "model_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\"\n",
    "model_file = model_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "\n",
    "output_dir = build_dir + \"/output_pre_and_post_proc\"\n",
    "\n",
    "#Delete previous run results if exist\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "    print(\"Previous run results deleted!\")\n",
    "\n",
    "build_steps = [\n",
    "    custom_step_add_pre_proc,\n",
    "    custom_step_add_post_proc,\n",
    "    \"step_qonnx_to_finn\",\n",
    "    \"step_tidy_up\",\n",
    "    \"step_streamline\",\n",
    "    \"step_convert_to_hw\",\n",
    "    \"step_create_dataflow_partition\",\n",
    "    \"step_specialize_layers\",\n",
    "    \"step_target_fps_parallelization\",\n",
    "    \"step_apply_folding_config\",\n",
    "    \"step_minimize_bit_width\",\n",
    "    \"step_generate_estimate_reports\",\n",
    "]\n",
    "\n",
    "cfg_estimates = build.DataflowBuildConfig(\n",
    "    output_dir          = output_dir,\n",
    "    mvau_wwidth_max     = 80,\n",
    "    target_fps          = 10000,\n",
    "    synth_clk_period_ns = 10.0,\n",
    "    fpga_part           = \"xc7z020clg400-1\",\n",
    "    steps               = build_steps,\n",
    "    generate_outputs=[\n",
    "        build_cfg.DataflowOutputType.ESTIMATE_REPORTS,\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0598b81",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "build.build_dataflow_cfg(model_file, cfg_estimates);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95230896",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -t -r {build_dir}/output_pre_and_post_proc/intermediate_models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a0263b1",
   "metadata": {},
   "source": [
    "You can use the code in the cell below to investigate the generated intermediate models. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44127417",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_to_investigate = \"custom_step_add_post_proc.onnx\"\n",
    "showInNetron(build_dir+\"/output_pre_and_post_proc/intermediate_models/\"+model_to_investigate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cc97505",
   "metadata": {},
   "source": [
    "Let's have a look at the model after the conversion to hw, to verify that now all layers are correctly converted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63131e3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron(build_dir+\"/output_pre_and_post_proc/intermediate_models/step_convert_to_hw.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fd0af6b",
   "metadata": {},
   "source": [
    "The model contains now a `Thresholding` layer in the beginning and a `LabelSelect` layer at the end. Please note, that there is still a `Transpose` node as the first layer of the graph, but we can solve this by converting the input data to the NHWC format before streaming it into the FINN accelerator."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6edf5c4-9213-45cd-834f-615c12685d9e",
   "metadata": {},
   "source": [
    "## Specialize layers configuration json <a id='specialize_layers'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ae83d6e-c704-4c7f-a922-a4b470c0a55f",
   "metadata": {},
   "source": [
    "The FINN compiler was developed with the assumption that the hardware blocks corresponding to the neural network layers are developed based on HLS. Although we do not want to abolish this HLS implementation at this time, it has become apparent over the years that for certain modules it makes sense to implement them in RTL. This allows us greater control over the resulting hardware and we can make optimal use of FPGA resources.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed72aabf-0517-422f-a686-6c70e7492114",
   "metadata": {},
   "source": [
    "So, with the growth of more and more RTL variants of common FINN hardware building blocks, we introduced an additional builder step called `step_specialize_layers`. In this step HW nodes get specialized to either an HLS or RTL variant of the node. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82a2bc39-8a37-49aa-a79d-2818e66ebd11",
   "metadata": {},
   "source": [
    "They get converted either based on pre-determined rules or the user provides a configuration file which contains the desired setting. If the user preference cannot be fulfilled, a warning will be printed and the implementation style will be set to a default. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc90b589-7a92-4996-9704-02736ac4e60e",
   "metadata": {},
   "source": [
    "The builder flow step before `step_create_dataflow_partition` generates a template json file to set the preferred implementation style per layer. We can copy it from one of the previous runs to this folder and manipulate it to pass it to a new build."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddb88eb1-3f11-4343-ae7c-3e5e8cbc34dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open(build_dir+\"/output_pre_and_post_proc/template_specialize_layers_config.json\", 'r') as json_file:\n",
    "    specialize_layers_config = json.load(json_file)\n",
    "\n",
    "print(json.dumps(specialize_layers_config, indent=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "158d7d8c-a072-4a50-9714-43ebaefa53d1",
   "metadata": {},
   "source": [
    "As you can see, each node is listed in the .json file and an empty string for the node attribute `preferred_impl_style` is instantiated by default. We can now use this .json and set the `preferred_impl_style` to pass to a new builder flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f464d35-6774-4751-80b4-b6230e501539",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(build_dir+\"/output_pre_and_post_proc/template_specialize_layers_config.json\", 'r') as json_file:\n",
    "    specialize_layers_config = json.load(json_file)\n",
    "\n",
    "# Set all preferred_impl_style to all HLS\n",
    "for key in specialize_layers_config:\n",
    "    if \"preferred_impl_style\" in specialize_layers_config[key]:\n",
    "        specialize_layers_config[key][\"preferred_impl_style\"] = \"hls\" \n",
    "# Save as .json    \n",
    "with open(\"specialize_layers_all_hls.json\", \"w\") as jsonFile:\n",
    "    json.dump(specialize_layers_config, jsonFile)\n",
    "         \n",
    "# Set SWG to RTL variant\n",
    "for key in specialize_layers_config:\n",
    "    if \"preferred_impl_style\" in specialize_layers_config[key]:\n",
    "        if key.startswith(\"ConvolutionInputGenerator\"):\n",
    "            specialize_layers_config[key][\"preferred_impl_style\"] = \"rtl\"\n",
    "        else:\n",
    "            specialize_layers_config[key][\"preferred_impl_style\"] = \"hls\"  \n",
    "# Save as .json    \n",
    "with open(\"specialize_layers_swg_rtl.json\", \"w\") as jsonFile:\n",
    "    json.dump(specialize_layers_config, jsonFile)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52592ea6-cd12-46b9-af91-5960b4749e7e",
   "metadata": {},
   "source": [
    "We created two `specialize_layers_config_files`:\n",
    "* One which sets all layers to `\"hls\"`\n",
    "* One that sets `preferred_impl_style` for the ConvolutionInputGenerator to `\"rtl\"`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "701905d8-c5cc-4cc0-b872-156c5b9d0432",
   "metadata": {},
   "source": [
    "In the following we will setup two build flows and run them to the estimate reports step. Afterwards we will investigate the intermediate .onnx files and compare the two runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22ff1a91-7ef7-44cb-86d3-60b9af7a8c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Build flow with custom folding configuration\n",
    "## specialize_layers_config_file = \"specialize_layers_all_hls.json\"\n",
    "\n",
    "model_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\"\n",
    "model_file = model_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "\n",
    "output_dir = build_dir + \"/output_all_hls\"\n",
    "\n",
    "#Delete previous run results if exist\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "    print(\"Previous run results deleted!\")\n",
    "\n",
    "build_steps = [\n",
    "    custom_step_add_pre_proc,\n",
    "    custom_step_add_post_proc,\n",
    "    \"step_qonnx_to_finn\",\n",
    "    \"step_tidy_up\",\n",
    "    \"step_streamline\",\n",
    "    \"step_convert_to_hw\",\n",
    "    \"step_create_dataflow_partition\",\n",
    "    \"step_specialize_layers\",\n",
    "    \"step_apply_folding_config\",\n",
    "    \"step_minimize_bit_width\",\n",
    "    \"step_generate_estimate_reports\",\n",
    "]\n",
    "\n",
    "cfg_estimates = build.DataflowBuildConfig(\n",
    "    output_dir          = output_dir,\n",
    "    mvau_wwidth_max     = 80,\n",
    "    synth_clk_period_ns = 10.0,\n",
    "    fpga_part           = \"xc7z020clg400-1\",\n",
    "    steps               = build_steps,\n",
    "    specialize_layers_config_file = \"specialize_layers_all_hls.json\",\n",
    "    generate_outputs=[\n",
    "        build_cfg.DataflowOutputType.ESTIMATE_REPORTS,\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9df41ff-ef6a-4d0e-ab36-241bb11ed241",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "build.build_dataflow_cfg(model_file, cfg_estimates);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff617f21-6001-4bb7-9cf7-2cc2acd3fbec",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Build flow with custom folding configuration\n",
    "## specialize_layers_config_file = \"specialize_layers_swg_rtl.json\"\n",
    "\n",
    "model_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\"\n",
    "model_file = model_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "\n",
    "output_dir = build_dir + \"/output_swg_rtl\"\n",
    "\n",
    "#Delete previous run results if exist\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "    print(\"Previous run results deleted!\")\n",
    "\n",
    "build_steps = [\n",
    "    custom_step_add_pre_proc,\n",
    "    custom_step_add_post_proc,\n",
    "    \"step_qonnx_to_finn\",\n",
    "    \"step_tidy_up\",\n",
    "    \"step_streamline\",\n",
    "    \"step_convert_to_hw\",\n",
    "    \"step_create_dataflow_partition\",\n",
    "    \"step_specialize_layers\",\n",
    "    \"step_apply_folding_config\",\n",
    "    \"step_minimize_bit_width\",\n",
    "    \"step_generate_estimate_reports\",\n",
    "]\n",
    "\n",
    "cfg_estimates = build.DataflowBuildConfig(\n",
    "    output_dir          = output_dir,\n",
    "    mvau_wwidth_max     = 80,\n",
    "    synth_clk_period_ns = 10.0,\n",
    "    fpga_part           = \"xc7z020clg400-1\",\n",
    "    steps               = build_steps,\n",
    "    specialize_layers_config_file = \"specialize_layers_swg_rtl.json\",\n",
    "    generate_outputs=[\n",
    "        build_cfg.DataflowOutputType.ESTIMATE_REPORTS,\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f48ba95-f7b5-455b-8041-25b7341ad115",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "build.build_dataflow_cfg(model_file, cfg_estimates);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bed4bedd-397d-4bd1-8531-c6ceac306715",
   "metadata": {},
   "source": [
    "First we are looking into the intermediate model after `step_create_dataflow_partition` and then after `step_specialize_layers`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e64db23-98cb-494b-851f-3cc2c3847451",
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron(build_dir+\"/output_all_hls/intermediate_models/step_create_dataflow_partition.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e1a6351-367f-47a6-b802-a2613ea455a1",
   "metadata": {},
   "source": [
    "Let's have a look first at the model which we specialize to \"all HLS\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f85d6c42-153d-4a40-b3cc-a4c8c89fe636",
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron(build_dir+\"/output_all_hls/intermediate_models/step_specialize_layers.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1520920-b7de-42a5-9ec8-e8503992fbd1",
   "metadata": {},
   "source": [
    "As you can see, each op type has now a suffix indicating that it is an HLS variant of the node. Additionally, when you click on one of the node in the Netron visualization, you can see that module is set to `finn.custom_op.fpgadataflow.hls`.\n",
    "\n",
    "Let's now have a look at the model in which we specialized the ConvolutionInputGenerator to `\"rtl\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f1f26a0-3a62-4920-bf40-5b1b798fa02e",
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron(build_dir+\"/output_swg_rtl/intermediate_models/step_specialize_layers.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f9c4de4-61ef-4698-ab23-87bf5953c5ae",
   "metadata": {},
   "source": [
    "You can use the cells above to try out different settings and pass it to the builder flow. Please note that not all layers have HLS and RTL variants, so it might be that the setting you define in `specialize_layers_config.json` gets ignored and a sensible default is set instead. The FINN compiler will display a warning in this case."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ffbadd1",
   "metadata": {},
   "source": [
    "## Folding configuration json <a id='folding_config'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c164040f",
   "metadata": {},
   "source": [
    "The FINN compiler allows the user to implement a network in streaming dataflow architecture, this means every layer is implemented individually and the data is streamed through the accelerator. We can customize each layer for specific performance and resource requirements by adjusting the parallelism and resource type of each layer. In the FINN context we refer to this customization of parallelism in each layer as folding. To learn more details about the influence of folding factors/parallelism in FINN, please have a look at our [folding tutorial](./3_folding.ipynb).\n",
    "\n",
    "In this section, we will look into the interface over which we can influence the customization of each layer using the FINN builder tool: A json file containing the folding configuration."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1299b86d",
   "metadata": {},
   "source": [
    "Depending on the invoked step, the FINN compiler can produce or consume a .json file containing the folding configuration for each layer. In the cell below, we will have a look at the automatically generated .json file, which is produced by `step_target_fps_parallelization`. We use this then as starting point to manipulate the folding configuration and feed it back into the builder tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f75f5634",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open(build_dir+\"/output_pre_and_post_proc/auto_folding_config.json\", 'r') as json_file:\n",
    "    folding_config = json.load(json_file)\n",
    "\n",
    "print(json.dumps(folding_config, indent=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8de787a7",
   "metadata": {},
   "source": [
    "As you can see from the printed cell above, the keys in the .json file are the node names of the layers in our network. For each of the layers, some node attributes are listed:\n",
    "* `PE` and `SIMD` are the folding parameters that determine the parallelism of each layer, depending on the layer they can be set to different values, for details refer to [this table](https://finn-dev.readthedocs.io/en/latest/internals.html#constraints-to-folding-factors-per-layer).\n",
    "* `mem_mode`: determines if the parameter memory will be implemented as part of the HLS/RTL code (`const`) or instantiated separately and connected with the layer over a memory streamer unit (`decoupled`). You can find more details in this part of the documentation: https://finn-dev.readthedocs.io/en/latest/internals.html#matrixvectoractivation-mem-mode . It is also possible to set the mem_mode to external which allows for the implementation for external weights.\n",
    "* `ram_style`: when selecting `decoupled` mode, the FINN compiler allows us to choose which memory resource will be used for the layer. The argument `ram_style` is set to the selected memory type:\n",
    "    * `auto`: Vivado will make the decision if the implementation is using LUTRAM or BRAM\n",
    "    * `distributed`: LUTRAM will be used\n",
    "    * `block`: BRAM will be used\n",
    "    * `ultra`: URAM will be used, if available on the selected board\n",
    "\n",
    "* `resType`: This is a node attribute for the MVAU layer and can be set to `lut` or `dsp`. Please note that selecting `dsp` will not enable the optimized RTL variant of the MVAU but rather generate HLS code utilizing DSPs, this is not optimal yet but can give an additional parameter for design space exploration.\n",
    "* `runtime_writeable_weights`: FINN offers the option to implement the weights as \"runtime writable\", this means you can write the weight values from the driver via an axilite interface."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd1519fe",
   "metadata": {},
   "source": [
    "In the following part of the tutorial, we will use the auto generated json file as starting point to create two new json files which explore the `ram_style` attribute. We will use one of the generated reports from the FINN builder to see the impact of these changes.\n",
    "For that, we will extract the total resources from the *estimate_layer_resources.json* report in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7f42774",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(build_dir+\"/output_pre_and_post_proc/report/estimate_layer_resources.json\", 'r') as json_file:\n",
    "    json_object = json.load(json_file)\n",
    "\n",
    "print(json.dumps(json_object[\"total\"], indent=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0be3b0e1",
   "metadata": {},
   "source": [
    "The FINN compiler estimates the network to use ~500 BRAM blocks and ~100k LUTs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4d177dc",
   "metadata": {},
   "source": [
    "We will use the `auto_folding_config.json` and create two folding configuration from that file:\n",
    "* All `ram_style` attributes set to `distributed`\n",
    "* All `ram_style` attributes set to `block`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "112af6fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(build_dir+\"/output_pre_and_post_proc/auto_folding_config.json\", 'r') as json_file:\n",
    "    folding_config = json.load(json_file)\n",
    "\n",
    "# Set all ram_style to LUT RAM\n",
    "for key in folding_config:\n",
    "    if \"ram_style\" in folding_config[key]:\n",
    "        folding_config[key][\"ram_style\"] = \"distributed\" \n",
    "# Save as .json    \n",
    "with open(\"folding_config_all_lutram.json\", \"w\") as jsonFile:\n",
    "    json.dump(folding_config, jsonFile)\n",
    "         \n",
    "# Set all ram_style to BRAM\n",
    "for key in folding_config:\n",
    "    if \"ram_style\" in folding_config[key]:\n",
    "        folding_config[key][\"ram_style\"] = \"block\" \n",
    "# Save as .json    \n",
    "with open(\"folding_config_all_bram.json\", \"w\") as jsonFile:\n",
    "    json.dump(folding_config, jsonFile)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e64a499",
   "metadata": {},
   "source": [
    "After generating these files, we will invoke the builder flow. To enable the FINN builder to take the generated folding configuration as input, we will need to set the additional builder argument `folding_config_file` and we will change the `build_steps` to not run `step_target_fps_parallelization`. The build step does not necessarily need to be excluded, but since we pass a separate folding configuration, the output from that step would be overwritten anyways, so we skip it for a faster execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdd9f706",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Build flow with custom folding configuration\n",
    "## folding_config_file = \"folding_config_all_lutram.json\"\n",
    "\n",
    "model_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\"\n",
    "model_file = model_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "\n",
    "output_dir = build_dir + \"/output_all_lutram\"\n",
    "\n",
    "#Delete previous run results if exist\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "    print(\"Previous run results deleted!\")\n",
    "\n",
    "build_steps = [\n",
    "    custom_step_add_pre_proc,\n",
    "    custom_step_add_post_proc,\n",
    "    \"step_qonnx_to_finn\",\n",
    "    \"step_tidy_up\",\n",
    "    \"step_streamline\",\n",
    "    \"step_convert_to_hw\",\n",
    "    \"step_create_dataflow_partition\",\n",
    "    \"step_specialize_layers\",\n",
    "    \"step_apply_folding_config\",\n",
    "    \"step_minimize_bit_width\",\n",
    "    \"step_generate_estimate_reports\",\n",
    "]\n",
    "\n",
    "cfg_estimates = build.DataflowBuildConfig(\n",
    "    output_dir          = output_dir,\n",
    "    mvau_wwidth_max     = 80,\n",
    "    synth_clk_period_ns = 10.0,\n",
    "    fpga_part           = \"xc7z020clg400-1\",\n",
    "    steps               = build_steps,\n",
    "    folding_config_file = \"folding_config_all_lutram.json\",\n",
    "    generate_outputs=[\n",
    "        build_cfg.DataflowOutputType.ESTIMATE_REPORTS,\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99b647c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "build.build_dataflow_cfg(model_file, cfg_estimates);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e705767d",
   "metadata": {},
   "source": [
    "We can now have a look at the produced model, when clicking on the individual nodes, you can see that all layers have the node attribute `ram_style` set to `distributed`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc680178",
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron(build_dir+\"/output_all_lutram/intermediate_models/step_generate_estimate_reports.onnx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "695ecfb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(build_dir+\"/output_all_lutram/report/estimate_layer_resources.json\", 'r') as json_file:\n",
    "    json_object = json.load(json_file)\n",
    "\n",
    "print(json.dumps(json_object[\"total\"], indent=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55208c70",
   "metadata": {},
   "source": [
    "The estimation report shows that BRAM utilization is down to zero and the LUT count went up to around 150k."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11b8430a",
   "metadata": {},
   "source": [
    "Let's do the same with the folding configuration which sets all memory resources to use BRAM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59e8aaaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Build flow with custom folding configuration\n",
    "## folding_config_file = \"folding_config_all_bram.json\"\n",
    "\n",
    "model_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\"\n",
    "model_file = model_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "\n",
    "output_dir = build_dir + \"/output_all_bram\"\n",
    "\n",
    "#Delete previous run results if exist\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "    print(\"Previous run results deleted!\")\n",
    "\n",
    "build_steps = [\n",
    "    custom_step_add_pre_proc,\n",
    "    custom_step_add_post_proc,\n",
    "    \"step_qonnx_to_finn\",\n",
    "    \"step_tidy_up\",\n",
    "    \"step_streamline\",\n",
    "    \"step_convert_to_hw\",\n",
    "    \"step_create_dataflow_partition\",\n",
    "    \"step_specialize_layers\",\n",
    "    \"step_apply_folding_config\",\n",
    "    \"step_minimize_bit_width\",\n",
    "    \"step_generate_estimate_reports\",\n",
    "]\n",
    "\n",
    "cfg_estimates = build.DataflowBuildConfig(\n",
    "    output_dir          = output_dir,\n",
    "    mvau_wwidth_max     = 80,\n",
    "    synth_clk_period_ns = 10.0,\n",
    "    fpga_part           = \"xc7z020clg400-1\",\n",
    "    steps               = build_steps,\n",
    "    folding_config_file = \"folding_config_all_bram.json\",\n",
    "    generate_outputs=[\n",
    "        build_cfg.DataflowOutputType.ESTIMATE_REPORTS,\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cdc1aa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "build.build_dataflow_cfg(model_file, cfg_estimates);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd0388fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "showInNetron(build_dir+\"/output_all_bram/intermediate_models/step_generate_estimate_reports.onnx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e60a3efb",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(build_dir+\"/output_all_bram/report/estimate_layer_resources.json\", 'r') as json_file:\n",
    "    json_object = json.load(json_file)\n",
    "\n",
    "print(json.dumps(json_object[\"total\"], indent=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97f87780",
   "metadata": {},
   "source": [
    "The initial implementation already had a high utilization of BRAM, but the estimations went now up to ~500 BRAMs while the LUT count went down to ~99k."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e65a8ded",
   "metadata": {},
   "source": [
    "You can use this example as a starting point to manipulate the folding configuration yourself. Instead of using the above code, you can also manually open one of the example .json files and set the values differently. Please be aware that the node attributes can not be set to arbitrary values. Especially the folding factors need to fulfil [certain constraints](https://finn-dev.readthedocs.io/en/latest/internals.html#constraints-to-folding-factors-per-layer). The other settings for node attributes, can be best looked up in the individual custom operator classes: [e.g. for MVAU](https://github.com/Xilinx/finn/blob/dev/src/finn/custom_op/fpgadataflow/matrixvectoractivation.py#L64)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a675834",
   "metadata": {},
   "source": [
    "## Additional builder arguments <a id='builder_arg'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7012b9a",
   "metadata": {},
   "source": [
    "In this section, we will have a peak into additional builder arguments the FINN compiler exposes. We will not be able to cover all but you will be able to have a look at a list and we encourage you to take your time to look into the different options there are to customize the FINN builder configuration."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "467d8829",
   "metadata": {},
   "source": [
    "We start by enabling the verification flow in the builder. The FINN compiler applies multiple transformations to the model before it gets turned into hardware, so we need to make sure that the functional behavior of the network does not change."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0c167f4",
   "metadata": {},
   "source": [
    "### Verification steps <a id='verify'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "308d52ba",
   "metadata": {},
   "source": [
    "Earlier in the tutorial, we had a look at how build steps are written. When investigating the `step_tidy_up`, we can see that before the changed model is returned a verification step can be run. In the case of `step_tidy_up` it is the step `\"initial python\"` that can be initiated by setting `VerificationStepType.TIDY_UP_PYTHON`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fe7318e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import finn.builder.build_dataflow_steps as build_dataflow_steps\n",
    "showSrc(build_dataflow_steps.step_tidy_up)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bbb84fb",
   "metadata": {},
   "source": [
    "Some of the default build steps have automatic verification enabled, when the corresponding verification step is set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce1aa025",
   "metadata": {},
   "outputs": [],
   "source": [
    "showSrc(build_cfg.VerificationStepType)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da1a2b88",
   "metadata": {},
   "source": [
    "In the cells below, we will use an example input from the CIFAR-10 data set and use the forward pass in Brevitas to generate a reference output. We save the input as `input.npy` and the reference output as `expected_output.npy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e157d03c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get golden io pair from Brevitas and save as .npy files\n",
    "from finn.util.test import get_trained_network_and_ishape, get_example_input, get_topk\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "(brevitas_model, ishape) = get_trained_network_and_ishape(\"cnv\", 2, 2)\n",
    "input_tensor_npy = get_example_input(\"cnv\")\n",
    "input_tensor_torch = torch.from_numpy(input_tensor_npy).float()\n",
    "input_tensor_torch = ToTensor().forward(input_tensor_torch).detach()\n",
    "output_tensor_npy = brevitas_model.forward(input_tensor_torch).detach().numpy()\n",
    "output_tensor_npy = get_topk(output_tensor_npy, k=1)\n",
    "\n",
    "np.save(\"input.npy\", input_tensor_npy)\n",
    "np.save(\"expected_output.npy\", output_tensor_npy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d03450e7",
   "metadata": {},
   "source": [
    "In the next step we set up the builder flow again, this time we will set the build argument `verify_steps` and pass a list of verification steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cd3032b",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Build flow with additional builder arguments enabled\n",
    "## verification steps\n",
    "\n",
    "model_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\"\n",
    "model_file = model_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "\n",
    "output_dir = build_dir + \"/output_with_verification\"\n",
    "\n",
    "#Delete previous run results if exist\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "    print(\"Previous run results deleted!\")\n",
    "\n",
    "build_steps = [\n",
    "    custom_step_add_pre_proc,\n",
    "    custom_step_add_post_proc,\n",
    "    \"step_qonnx_to_finn\",\n",
    "    \"step_tidy_up\",\n",
    "    \"step_streamline\",\n",
    "    \"step_convert_to_hw\",\n",
    "    \"step_create_dataflow_partition\",\n",
    "    \"step_specialize_layers\",\n",
    "    \"step_target_fps_parallelization\",\n",
    "    \"step_apply_folding_config\",\n",
    "    \"step_minimize_bit_width\",\n",
    "    \"step_generate_estimate_reports\",\n",
    "]\n",
    "\n",
    "cfg_estimates = build.DataflowBuildConfig(\n",
    "    output_dir          = output_dir,\n",
    "    mvau_wwidth_max     = 80,\n",
    "    target_fps          = 10000,\n",
    "    synth_clk_period_ns = 10.0,\n",
    "    fpga_part           = \"xc7z020clg400-1\",\n",
    "    steps               = build_steps,\n",
    "    generate_outputs=[\n",
    "        build_cfg.DataflowOutputType.ESTIMATE_REPORTS,\n",
    "    ],\n",
    "    verify_steps=[\n",
    "        build_cfg.VerificationStepType.QONNX_TO_FINN_PYTHON,\n",
    "        build_cfg.VerificationStepType.TIDY_UP_PYTHON,\n",
    "        build_cfg.VerificationStepType.STREAMLINED_PYTHON,\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d05b985",
   "metadata": {},
   "source": [
    "When execution the code below, the verification will be invoked in the background. After the execution we can check if the verification was successful by investigating the output directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3a46e76",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "build.build_dataflow_cfg(model_file, cfg_estimates);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca1d571d",
   "metadata": {},
   "source": [
    "The output directory has now an additional directory called `verification_output`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca74d537",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls {build_dir}/output_with_verification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "908ecda4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls {build_dir}/output_with_verification/verification_output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcbc6f49",
   "metadata": {},
   "source": [
    "The directory contains three .npy files. These files are the saved output files from the different verification steps. The suffix indicates if the array matches with the expected output. In our case, the suffix is for all verification steps `_SUCCESS`. Since the outputs are saved as .npy, we can open and investigate the files simply in Python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a1b6ca9",
   "metadata": {},
   "outputs": [],
   "source": [
    "verify_initial_python = np.load(build_dir + \"/output_with_verification/verification_output/verify_initial_python_0_SUCCESS.npy\")\n",
    "print(\"The output of the verification step after the step_tidy_up is: \" + str(verify_initial_python))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6558e19e",
   "metadata": {},
   "source": [
    "If the generated output does not match the expected output, these files can be used for debugging."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4609f94d",
   "metadata": {},
   "source": [
    "### Other builder arguments <a id='other_args'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37b6853d",
   "metadata": {},
   "source": [
    "Next to the enablement of the verification flows, the FINN builder has numerous additional builder arguments to further customize your network. \n",
    "Let's have a look at the options for the arguments. We want to only filter out the FINN specific arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9f6aa29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter out methods\n",
    "builder_args = [m for m in dir(build_cfg.DataflowBuildConfig) if not m.startswith('_')]\n",
    "print(\"\\n\".join(builder_args))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b12ab370",
   "metadata": {},
   "source": [
    "There are attributes that come from the dataclasses-json class: `to_dict`, `to_json`, `schema`, `from_json`, `from_dict`. This class is used for the implementation of the FINN builder. In this tutorial, we are mainly interested in the FINN specific arguments.  \n",
    "\n",
    "Some of these arguments we have seen already in the Cybersecurity notebook and in this notebook, e.g. `target_fps`, `fpga_part` and `folding_config_file`. In the code of the FINN builder, the function of each builder argument is documents, you can have a look [here](https://github.com/Xilinx/finn/blob/dev/src/finn/builder/build_dataflow_config.py#L155) and scroll through the available builder arguments."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9aba0493",
   "metadata": {},
   "source": [
    "So far, in this notebook, we only looked at configurations up to the generation of estimate reports, a lot of these builder arguments actually become relevant at a later stage in the FINN flow.\n",
    "\n",
    "Let's have a look at the default build dataflow steps for the complete FINN flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec39b9f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\n\".join(build_cfg.default_build_dataflow_steps))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9bc5715",
   "metadata": {},
   "source": [
    "You can see that after the generation of the estimate reports, the code generation and the ip generation is invoked (`step_hw_codegen` and `step_hw_ipgen`). The FIFO depths are determined and the FIFOs are inserted in the network (`step_set_fifo_depths`), we can then create an IP design of our whole network by stitching the IPs from each layer together (`step_create_stitched_ip`). At this point we have an implementation of the neural network that we can integrate within a bigger FPGA design, we can run performance measurements using simulation (`step_measure_rtlsim_performance`) and out-of-context synthesis (`step_out_of_context_synthesis`) for it.\n",
    "The FINN builder also provides automatic system integration for Zynq and Alveo devices, this can be invoked by running `step_synthesize_bitfile`, `step_make_pynq_driver` and `step_deployment_package`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76df000f",
   "metadata": {},
   "source": [
    "You can have a closer look at each step by either using the `showSrc()` function or by accessing the doc string."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caf49f03",
   "metadata": {},
   "outputs": [],
   "source": [
    "import finn.builder.build_dataflow_steps as build_dataflow_steps\n",
    "print(build_dataflow_steps.step_hw_codegen.__doc__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c84a9fbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "showSrc(build_dataflow_steps.step_hw_codegen)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c249f141",
   "metadata": {},
   "source": [
    "This concludes the advanced builder settings tutorial. Below you can find code that can help you investigating more of the builder arguments and invoking the whole flow to generate a bitfile."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b98eb65",
   "metadata": {},
   "source": [
    "### Example for additional builder arguments & bitfile generation <a id='example_args'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dbdab42",
   "metadata": {},
   "source": [
    "#### Standalone Thresholds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e21ff36f",
   "metadata": {},
   "source": [
    "In FINN, convolutions are expressed with three components:\n",
    "* An Im2Col operation\n",
    "* A matrix multiplication\n",
    "* A MultiThreshold operation\n",
    "\n",
    "When converting these nodes into HW layers, by default the MatMul and the MultiThreshold gets converted into **one** component called Matrix-Vector-Activation Unit (MVAU). But the FINN compiler allows us to implement the activation separately. This gives an additional possibility for customization because we can adjust the folding parameters of the standalone threshold unit independently. \n",
    "\n",
    "If you would like to enable this feature, you can set the build argument `standalone_thresholds` to `True`. In the code below this feature is enabled and you can have a look at the generated .onnx file. Please note that you need to uncomment the code first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2619ebde",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Build flow with additional builder arguments enabled\n",
    "## standalone_thresholds = True\n",
    "\n",
    "model_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\"\n",
    "model_file = model_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "\n",
    "output_dir = build_dir + \"/output_standalone_thresholds\"\n",
    "\n",
    "#Delete previous run results if exist\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "    print(\"Previous run results deleted!\")\n",
    "\n",
    "build_steps = [\n",
    "    custom_step_add_pre_proc,\n",
    "    custom_step_add_post_proc,\n",
    "    \"step_qonnx_to_finn\",\n",
    "    \"step_tidy_up\",\n",
    "    \"step_streamline\",\n",
    "    \"step_convert_to_hw\",\n",
    "    \"step_create_dataflow_partition\",\n",
    "    \"step_specialize_layers\",\n",
    "    \"step_target_fps_parallelization\",\n",
    "    \"step_apply_folding_config\",\n",
    "    \"step_minimize_bit_width\",\n",
    "    \"step_generate_estimate_reports\",\n",
    "]\n",
    "\n",
    "cfg_estimates = build.DataflowBuildConfig(\n",
    "    output_dir            = output_dir,\n",
    "    mvau_wwidth_max       = 80,\n",
    "    target_fps            = 10000,\n",
    "    synth_clk_period_ns   = 10.0,\n",
    "    fpga_part             = \"xc7z020clg400-1\",\n",
    "    standalone_thresholds = True,\n",
    "    steps                 = build_steps,\n",
    "    generate_outputs=[\n",
    "        build_cfg.DataflowOutputType.ESTIMATE_REPORTS,\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2e9bc42",
   "metadata": {},
   "outputs": [],
   "source": [
    "#%%time\n",
    "#build.build_dataflow_cfg(model_file, cfg_estimates);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32ae296e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#showInNetron(build_dir+\"/output_standalone_thresholds/intermediate_models/step_generate_estimate_reports.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "601eb5f8",
   "metadata": {},
   "source": [
    "#### Run the whole flow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42aa929b",
   "metadata": {},
   "source": [
    "The code below can be used to invoke the full builder flow and obtain more output products, be aware that this runs synthesis and bitfile generation and it might take over an hour. Please note that you need to uncomment the code first."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffa2a352",
   "metadata": {},
   "source": [
    "For an optimized design, we saved a local copy of the folding configuration for cnv-w2a2 on the Pynq-Z1 board from [finn-examples](https://github.com/Xilinx/finn-examples) in this folder. And will pass it to the build flow. Please also note below that we now pass the board as argument to the builder (`board = \"Pynq-Z1\"`) instead of just the fpga part. This time we will select all possible outputs to generate. Please be aware that running the full build might take a few hours."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d1b041f-027c-444e-81ac-98ce9b6d1b51",
   "metadata": {},
   "source": [
    "Note that we set one additional argument: `default_swg_exception = True`. This is done because this example is customized to fit on the Pynq-Z1 board, to optimize the resources we remove FIFOs between SWGs and MVAUs manually to avoid unnecessary buffering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4efd46f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import finn.builder.build_dataflow as build\n",
    "import finn.builder.build_dataflow_config as build_cfg\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "## Build flow with hardware build\n",
    "\n",
    "model_dir = os.environ['FINN_ROOT'] + \"/notebooks/advanced\"\n",
    "model_file = model_dir + \"/end2end_cnv_w2a2_export.onnx\"\n",
    "\n",
    "output_dir = build_dir + \"/output_bitfile\"\n",
    "\n",
    "#Delete previous run results if exist\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "    print(\"Previous run results deleted!\")\n",
    "\n",
    "build_steps = [\n",
    "    custom_step_add_pre_proc,\n",
    "    custom_step_add_post_proc,\n",
    "    \"step_qonnx_to_finn\",\n",
    "    \"step_tidy_up\",\n",
    "    \"step_streamline\",\n",
    "    \"step_convert_to_hw\",\n",
    "    \"step_create_dataflow_partition\",\n",
    "    \"step_specialize_layers\",\n",
    "    \"step_target_fps_parallelization\",\n",
    "    \"step_apply_folding_config\",\n",
    "    \"step_minimize_bit_width\",\n",
    "    \"step_generate_estimate_reports\",\n",
    "    \"step_hw_codegen\",\n",
    "    \"step_hw_ipgen\",\n",
    "    \"step_set_fifo_depths\",\n",
    "    \"step_create_stitched_ip\",\n",
    "    \"step_measure_rtlsim_performance\",\n",
    "    \"step_out_of_context_synthesis\",\n",
    "    \"step_synthesize_bitfile\",\n",
    "    \"step_make_pynq_driver\",\n",
    "    \"step_deployment_package\",\n",
    "]\n",
    "\n",
    "cfg_build = build.DataflowBuildConfig(\n",
    "    output_dir                    = output_dir,\n",
    "    mvau_wwidth_max               = 80,\n",
    "    synth_clk_period_ns           = 10.0,\n",
    "    #specialize_layers_config_file = \"specialize_layers_all_hls.json\",\n",
    "    folding_config_file           = \"cnv-w2a2_folding_config.json\",\n",
    "    board                         = \"Pynq-Z1\",\n",
    "    shell_flow_type               = build_cfg.ShellFlowType.VIVADO_ZYNQ,\n",
    "    steps                         = build_steps,\n",
    "    default_swg_exception         = True,\n",
    "    generate_outputs=[\n",
    "        build_cfg.DataflowOutputType.ESTIMATE_REPORTS,\n",
    "        build_cfg.DataflowOutputType.STITCHED_IP,\n",
    "        build_cfg.DataflowOutputType.RTLSIM_PERFORMANCE,\n",
    "        build_cfg.DataflowOutputType.OOC_SYNTH,\n",
    "        build_cfg.DataflowOutputType.BITFILE,\n",
    "        build_cfg.DataflowOutputType.PYNQ_DRIVER,\n",
    "        build_cfg.DataflowOutputType.DEPLOYMENT_PACKAGE,\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7ff6c19",
   "metadata": {},
   "outputs": [],
   "source": [
    "#%%time\n",
    "#build.build_dataflow_cfg(model_file, cfg_build);"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
